import cv2
import numpy as np
import argparse
from utils import FPS


trackers = cv2.MultiTracker_create()
target = 0

# opencv已经实现了的追踪算法

OPENCV_OBJECT_TRACKERS = {
	"csrt": cv2.TrackerCSRT_create,
	"kcf": cv2.TrackerKCF_create,
	"boosting": cv2.TrackerBoosting_create,
	"mil": cv2.TrackerMIL_create,
	"tld": cv2.TrackerTLD_create,
	"medianflow": cv2.TrackerMedianFlow_create,
	"mosse": cv2.TrackerMOSSE_create
}

# 初始化参数
confThreshold = 0.3  # 置信度阈值
nmsThreshold = 0.3  # 非最大抑制阈值
inpWidth = 416  # 网络输入图像的宽度
inpHeight = 416  # 网络输入图像的高度

# 加载类名
classesFile = 'coco.names.txt'
classes = None
with open(classesFile, 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')


# 模型的配置文件和权值文件
modelConfiguration = "yolov3/yolov3.cfg"
modelWeights = "yolov3/yolov3.weights"

print("[INFO] loading model...")
net = cv2.dnn.readNetFromDarknet(modelConfiguration, modelWeights)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)


# 获取输出层的名称
def getOutputsNames(net):
    # 获取网络中所有层的名称
    layersNames = net.getLayerNames()
    # 获取输出层的名称
    return [layersNames[i[0] - 1] for i in net.getUnconnectedOutLayers()]

# 绘制预测得到的边界框
def drawPred(boxes,classIds):
    for box in boxes:
        (x, y, w, h) = [int(v) for v in box]
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 255), 2)

    for box, id in zip(boxes, classIds):
        cv2.putText(frame, classes[id], (int(box[0]), int(box[1]) - 2),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 255, 255), 2)

# 使用非最大值抑制移除低置信度的边界框
def postprocess(frame, outs):
    frameHeight = frame.shape[0]
    frameWidth = frame.shape[1]

    # 扫描从网络输出的所有边界框并仅保留
    # 置信度得分很高的边界框，将框的类标签指定为具有最高分数的类。
    classIds = []
    confidences = []
    boxes = []
    for out in outs:
        for detection in out:
            scores = detection[5:]
            classId = np.argmax(scores)
            confidence = scores[classId]
            if confidence > confThreshold:
                center_x = int(detection[0] * frameWidth)
                center_y = int(detection[1] * frameHeight)
                width = int(detection[2] * frameWidth)
                height = int(detection[3] * frameHeight)
                left = int(center_x - width / 2)
                top = int(center_y - height / 2)
                if classes[classId] != "person":
                    continue
                classIds.append(classId)
                confidences.append(float(confidence))
                boxes.append([left, top, width, height])

    # 执行非最大抑制以消除置信度较低的冗余重叠框
    indices = cv2.dnn.NMSBoxes(boxes, confidences, confThreshold, nmsThreshold)
    global target
    target = len(indices)
    for i in indices:
        i = i[0]
        box = boxes[i]
        left = box[0]
        top = box[1]
        width = box[2]
        height = box[3]
        # drawPred(classIds[i], confidences[i], left, top, left + width, top + height)
    return classIds,boxes
# 处理输入
winName = 'yolov3 detect and opencv_track'
cv2.namedWindow(winName, cv2.WINDOW_NORMAL)

outputFile = "out"
print("[INFO] starting video stream...")
cap = cv2.VideoCapture("vtest.avi")

vid_writer = cv2.VideoWriter(outputFile, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 30,
                             (round(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), round(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))

# 计算FPS
fps = FPS().start()


while cv2.waitKey(1) < 0:
    # 从视频获取帧
    hasFrame, frame = cap.read()

    # 在视频结束时终止程序
    if not hasFrame:
        print('Done processing !!!')
        print('Output file is stored as ', outputFile)
        cv2.waitKey(3000)
        break

    # 从框架创建4D blob。
    blob = cv2.dnn.blobFromImage(frame, 1 / 255, (inpWidth, inpHeight), [0, 0, 0], 1, crop=False)

    # 设置网络的输入
    net.setInput(blob)

    # 运行前向传递以获得输出层的输出
    outs = net.forward(getOutputsNames(net))

    # 移除低置信度的边界框
    classIds, boxes = postprocess(frame, outs)

    drawPred(boxes, classIds)

    vid_writer.write(frame.astype(np.uint8))
    cv2.imshow(winName, frame)

    # 计算FPS
    fps.update()

fps.stop()
print("[INFO] elapsed time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))
cap.release()